import json
import time
from typing import AsyncGenerator, Dict, List, Optional

import httpx

from src.const import CHUNK_TYPE, MODEL_MAPPING
from src.schemas.chat import ChatCompletionChunk, Choice, ChoiceDelta, Message


def get_model_info(model_name: str) -> Optional[Dict]:
    return MODEL_MAPPING.get(model_name.lower(), None)


def parse_messages(messages: List[Message]) -> str:
    only_user_message = True
    for m in messages:
        if m.role == "user":
            only_user_message = False
            break
    if only_user_message:
        prompt = "\n".join([f"{m.role}: {m.content}" for m in messages])
    else:
        prompt = "\n".join([f"{m.content}" for m in messages])
    return prompt


async def process_response_stream(response: httpx.Response, model_id: str) -> AsyncGenerator[str, None]:
    """处理流式响应，带有超时控制和错误处理"""
    def _create_chunk(content: str = "", reasoning_content: str = "", finish_reason: Optional[str] = None) -> str:
        choice_delta = ChoiceDelta(content=content, reasoning_content=reasoning_content)
        choice = Choice(delta=choice_delta, finish_reason=finish_reason)
        chunk = ChatCompletionChunk(created=int(time.time()), model=model_id, choices=[choice])
        return chunk.model_dump_json(exclude_unset=True)

    last_activity = time.time()
    STREAM_TIMEOUT = 30  # 流式响应超时时间（秒）
    status = ""
    start_word = "data: "

    try:
        async for line in response.aiter_lines():
            # 检查流式响应是否超时
            current_time = time.time()
            if current_time - last_activity > STREAM_TIMEOUT:
                yield _create_chunk("", "", "timeout")
                return
            last_activity = current_time

            # 去除 line 开头的 start_word
            # if line and line.startswith(start_word):
            #     line = line[len(start_word):]
            # if not line:
            #     continue
            # yield line
            # continue

            # 处理数据行
            if not line or not line.startswith(start_word):
                continue
            
            data: str = line[len(start_word):]
            
            # 处理特殊消息
            if data == "[DONE]":
                yield "[DONE]"
                return
            
            # 处理状态更新
            if data in (CHUNK_TYPE.STATUS, CHUNK_TYPE.SEARCH_WITH_TEXT, CHUNK_TYPE.REASONER, CHUNK_TYPE.TEXT):
                status = data
                continue
            
            # 跳过非JSON数据
            if not data.startswith("{"):
                continue

            try:
                chunk_data: Dict = json.loads(data)
                
                # 处理不同类型的消息
                if chunk_data.get("content"):
                    yield _create_chunk("", chunk_data["content"])
                if chunk_data.get("msg"):
                    yield _create_chunk(chunk_data["msg"])
                # continue

            except json.JSONDecodeError:
                continue  # 忽略无效的JSON数据
                
    except httpx.ReadTimeout:
        yield _create_chunk("", "", "timeout")
    except Exception as e:
        yield _create_chunk(f"Error: {str(e)}", "", "error")
    finally:
        # 确保响应被正确关闭
        if not response.is_closed:
            await response.aclose()
